import re
from functools import cached_property
from typing import Any, Generator
from typing_extensions import Self
from pydantic import BaseModel
from pysam import VariantRecord  # pylint: disable=E0611
from .enum import LoF


class Annotation(BaseModel):
    gene: str
    transcript: str
    exon: str
    na_change: str
    aa_change: str

    @classmethod
    def from_detail(cls, detail) -> Self:
        return cls.model_validate(dict(zip("gene", "transcript", "exon", "na_change", "aa_change"), detail.split(":")[:5]))


class SNV(BaseModel):
    # pk: str
    chrom: str
    pos: int
    ref: str
    alt: str
    gene: str
    transcript: str
    event: str
    na_change: str
    aa_change: str

    class Config:
        arbitrary_types_allowed = True

    @cached_property
    def lof(self) -> LoF | None:
        return LoF.from_event(self.event)

    @cached_property
    def termination(self):
        """
        Get termination position from pHGVS:
        # NM_031475.2:p.Gln98*
        # NM_031475.2:p.Ala586Glyfs*73
        # NP_000305.3:p.Arg378SerfsTer5
        # p.Arg97Glyfs*26 (alternatively p.Arg97GlyfsTer26, or short p.Arg97fs)
        """
        m = re.compile(r"p\.\D+(\d+)\D+fs(\*|X|Ter)(\d+)").search(self.aa_change)
        if m:
            return int(m.group(1)) + int(m.group(3))
        m = re.compile(r"p\.\D+(\d+)fs").search(self.aa_change)
        if m:
            return int(m.group(1))
        m = re.compile(r"p\.\D+(\d+)(\*|X|Ter)").search(self.aa_change)
        if m:
            return int(m.group(1))
        return -1

    @classmethod
    def from_variant_record(cls, record: VariantRecord) -> Generator[list[Self], Any, None]:
        if record.info.get("EVENT"):
            for event, detail in filter(lambda x: "." not in x, zip(record.info["EVENT"], record.info["DETAIL"])):
                snvs = []
                for d in map(lambda x: x.split(":"), detail.split("|")):
                    snvs.append(
                        SNV(
                            chrom=record.chrom,
                            pos=record.pos,
                            ref=record.ref,
                            alt=record.alts[0],
                            event=event,
                            gene=d[0],
                            transcript=d[1].split(".", 1)[0],
                            na_change=d[3],
                            aa_change=d[4] if len(d) == 5 else "",
                        )
                    )
                yield snvs
